import json5
import sys
import argparse
from flow import get_flow
from utils.parse_input import parse_input

def main():
    """
    Main function for the OnceOne Agent Core.
    
    Supports two command modes:
    1. chat: Free Conversation Mode for natural language dialogue
    2. check: Correction Mode for analyzing and correcting JSON input
    
    Usage:
        python main.py chat  # Enter conversation mode
        python main.py check  # Enter correction mode
    """
    # Parse command line arguments
    parser = argparse.ArgumentParser(description='OnceOne Agent Core')
    subparsers = parser.add_subparsers(dest='mode', help='Operation mode')
    
    # Add chat mode subparser
    chat_parser = subparsers.add_parser('chat', help='Enter conversation mode')
    
    # Add check mode subparser
    check_parser = subparsers.add_parser('check', help='Enter correction mode')
    
    args = parser.parse_args()
    
    # If no mode is specified, show help and exit
    if not args.mode:
        parser.print_help()
        sys.exit(1)
    
    # Enter the appropriate mode
    if args.mode == 'chat':
        chat_mode()
    elif args.mode == 'check':
        check_mode()

def chat_mode():
    """Handle chat (free conversation) mode with multi-round dialogue."""
    print("\nEntering chat mode. Type 'exit' to quit.\n")
    
    # Initialize conversation history
    conversation_history = []
    
    while True:
        # Get user input
        print("> ", end="")
        user_input = input()
        
        # Check for exit command
        if user_input.lower() == 'exit':
            print("\nExiting chat mode. Goodbye!")
            break
        
        # Initialize shared store for this conversation turn
        shared = {
            "input_type": "chat",
            "user_input": None,
            "conversation_history": conversation_history,
            "response": None
        }
        
        try:
            # Parse input for chat mode
            parsed_input = parse_input(user_input, "chat")
            shared["user_input"] = parsed_input
            
            # Get the appropriate flow for chat mode and run it
            flow = get_flow("chat")
            flow.run(shared)
            
            # Update conversation history
            conversation_history = shared["conversation_history"]
            
            # Output the response
            print(f"\n{shared['response']}\n")
            
        except Exception as e:
            print(f"\nError: {e}\n")

def check_mode():
    """Handle check (correction) mode with multi-round dialogue, only accepting valid JSON input."""
    print("\nEntering check mode. Type 'exit' to quit.\n")
    print("Note: In check mode, only valid JSON input is accepted.\n")
    
    while True:
        # Get user input
        print("> ", end="")
        user_input = input()
        
        # Check for exit command
        if user_input.lower() == 'exit':
            print("\nExiting check mode. Goodbye!")
            break
        
        # Initialize shared store for this correction turn
        shared = {
            "input_type": "check",
            "user_input": None,
            "response": None
        }
        
        try:
            # Parse input for check mode
            parsed_input = parse_input(user_input, "check")
            shared["user_input"] = parsed_input
            
            # Get the appropriate flow for check mode and run it
            flow = get_flow("check")
            flow.run(shared)
            
            # Output the response
            print(f"\n{shared['response']}\n")
            
        except ValueError as e:
            print(f"\nError: {e}\nPlease provide valid JSON input according to the required format.\n")
        except Exception as e:
            print(f"\nError: {e}\n")

if __name__ == "__main__":
    main()
